#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
MUD客户端网络协议处理模块
负责处理与MUD服务器的协议交互，包括Telnet协议选项协商等
"""
import zlib
import logging
import telnetlib3
import datetime
import sys
import zlib
from .steam_write import MUDTelnetWriter, MUDTelnetWriterUnicode


class MUDTelnetReaderUnicode(telnetlib3.TelnetReaderUnicode):
    """
    MUD客户端的Telnet读取器Unicode类
    继承自TelnetReaderUnicode，用于处理Unicode字符的读取
    """
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)


class MUDTelnetClient(telnetlib3.TelnetClient):
    """
    MUD Telnet客户端
    继承自BaseClient，处理与MUD服务器的协议交互，包括Telnet选项协商、MCCP压缩等
    """
    
    _reader_factory_encoding = MUDTelnetReaderUnicode
    _writer_factory = MUDTelnetWriter
    _writer_factory_encoding = MUDTelnetWriterUnicode

    DEFAULT_LOCALE = "zh_CN"


    def __init__(self, *args, **kwargs):
        """
        初始化Telnet客户端
        
        Args:
            term (str): 终端类型，默认为"XTERM-256COLOR"
        """
        super().__init__(*args, **kwargs)

        self.log = logging.getLogger(__name__)
        self.log.setLevel(logging.DEBUG)

        # 初始化MCCP2相关属性
        self.decompressor = zlib.decompressobj(-15)  # 使用原始deflate算法
        self.mccp_enabled = False  # 标记MCCP是否已启用
        self.mccp_buffer = bytearray()  # 用于存储MCCP压缩数据的缓冲区
        
        # 初始化MCCP2相关属性
        self.mccp2_compressing = False
        # 初始化zlib解压缩器，使用wbits=-15处理MCCP2压缩数据（无zlib头）
        self.decompressor = zlib.decompressobj()
        # 标志位，用于标记是否是第一次处理压缩数据
        self.mccp2_first_run = True
        
    def data_received(self, data):
        """
        处理原始数据
        
        Args:
            data (bytes): 接收到的数据
        """
        # 检查是否启用了MCCP2压缩
        if self.mccp2_compressing and self.mccp2_first_run:
            # 如果打开且是第一次运行，则动态查找MCCP2子协商结束位置
            # MCCP2子协商命令序列: IAC SB MCCP2 IAC SE (\xff \xfa V \xff \xf0)
            mccp2_end_sequence = b'\xff\xfaV\xff\xf0'
            seq_pos = data.find(mccp2_end_sequence)
            
            if seq_pos != -1:
                # 找到了MCCP2子协商结束序列，确定压缩数据开始位置
                compressed_start = seq_pos + len(mccp2_end_sequence)
                plain_data = data[:compressed_start]
                compressed_data = data[compressed_start:]
                
                # 提交普通数据
                super().data_received(plain_data)
                
                # 对剩余数据进行解压缩
                self._handle_mccp2_data(compressed_data)
                self.mccp2_first_run = False
            else:
                # 没有找到MCCP2子协商结束序列，全部作为普通数据处理
                super().data_received(data)
        elif self.mccp2_compressing and not self.mccp2_first_run:
            # 如果打开但不是第一次运行，则进行数据处理
            self._handle_mccp2_data(data)
        else:
            # 未启用MCCP2压缩，直接处理数据
            super().data_received(data)

    def _handle_mccp2_data(self, data):
        """
        处理MCCP2数据
        
        Args:
            data (bytes): 接收到的数据
        """
        # 尝试进行解压缩
        try:
            decompressed_data = self.decompressor.decompress(data)
            if decompressed_data:
                super().data_received(decompressed_data)        
        except zlib.error as e:
            self.log.warning(f"MCCP2解压缩错误: {e}")
            # 出错时关闭压缩
            self.mccp2_compressing = False
            self.mccp2_first_run = False
            self.decompressor = None
            
            # 将原始数据提交给父类处理
            super().data_received(data)
        except Exception as e:
            self.log.error(f"处理MCCP2数据时出错: {e}")
            # 将原始数据提交给父类处理
            super().data_received(data)

    def _handle_sb_mccp2(self, buf):
        """
        处理MCCP2子协商数据
        
        Args:
            buf (collections.deque): MCCP2子协商数据缓冲区
        """
        # 设置压缩标志
        self.mccp2_compressing = True
        # 重置第一次运行标志
        self.mccp2_first_run = True
        self.log.debug("MCCP2压缩已启用")

    def encoding(self, outgoing=False, incoming=False):
        """
        重写encoding方法，强制返回utf-8编码
        
        Args:
            outgoing (bool): 是否用于输出方向
            incoming (bool): 是否用于输入方向
            
        Returns:
            str: 编码名称
        """
        # 检查参数是否正确设置
        if not (outgoing or incoming):
            raise TypeError(
                "encoding arguments 'outgoing' and 'incoming' "
                "are required: toggle at least one."
            )
            
        # 只有在force_binary或BINARY选项协商通过时，TelnetClient.encoding()才会返回正确的编码
        # 为了避免这个问题，我们直接返回，初始化时default_encoding=encoding参数值
        return self.default_encoding

    def get_cached_data(self):
        """
        从transport中把已缓存的压缩数据拿出来
        
        Returns:
            bytes: 已缓存的原始字节数据
        """
        transport = self._transport
        if not transport:
            return b''

         # 添加异常处理，以防transport不支持暂停操作
        try:
            transport.pause_reading()
        except NotImplementedError:
            # 如果transport不支持暂停，则继续执行但记录日志
            self.log.warning("Transport does not support pause_reading()")

        # 把内核/transport 里已缓存的 raw 字节拿出来
        raw = bytes(getattr(transport, '_buffer', b''))
        transport._buffer = bytearray()

        # 保存原来的"文本解析"入口
        self._original_data_received = self.data_received
        # 从此所有 raw 字节先走 zlib
        self.data_received = self.mccp_received

            # 恢复读取操作
        try:
            transport.resume_reading()
        except (NotImplementedError, AttributeError):
            # 如果transport不支持恢复读取或transport为None，则继续执行
            pass
        self.log.debug(f"取回缓存数据: {len(raw)}字节")

        return raw
    
    def mccp_received(self, data):
        """
        以后所有 TCP 字节都先到这，解压后再走正常 telnet 解析
        """
        if self.decompressor is None:
            self.decompressor = zlib.decompressobj(-15)  # 使用原始deflate算法
            
        try:
            self.log.debug(f"MCCP解压: {len(data)}字节")
            plain = self.decompressor.decompress(data)
            if plain:
                # 原路丢给原来的解析器，上层毫无感知
                self._original_data_received(plain)
        except zlib.error as e:
            self.log.error(f"MCCP解压错误: {e}")
            # 解压失败时重置解压器
            self.decompressor = zlib.decompressobj(-15)
            # 发送DONT MCCP2命令通知服务器停止压缩
            self.writer.iac(telnetlib3.DONT, b'V')
            # 将原始数据传递给原始处理方法
            self._original_data_received(data)